package tensorflow.image.classify;

/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

import org.tensorflow.*;

import java.io.File;
import java.io.FileFilter;
import java.io.IOException;
import java.nio.charset.Charset;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Iterator;
import java.util.List;

/**
 * Simplified version of
 * https://github.com/tensorflow/tensorflow/blob/r1.4/tensorflow/java/src/main/java/org/tensorflow/examples/LabelImage.java
 *
 * 图像分类
 * java调用已经生成好的tensorflow模型进行分类
 */
public class ImageClassify {

  private static String dataPath = "/data/model";

  public static void main(String[] args) throws Exception {

    // 数据label
    final List<String> labels = loadLabels(dataPath + "/labels.txt");

    try (Graph graph = new Graph();
         Session session = new Session(graph)) {

      // 导入预测模型
      graph.importGraphDef(loadTensorGraph(dataPath + "/graph.pb"));

      // 用于预测的图片
      File files = new File(dataPath + "/validate_image");
      File[] validate_files = files.listFiles(new FileFilter() {
        @Override
        public boolean accept(File file) {
          if (file.isFile()) {
            return true;
          }
          return false;
        }
      });

      //printAllOperation(graph);

      /**
       * 预测结果是一个1*labels.size()的二维数组
       *
       */
      float[][] predictionData = null;
      for (File f : validate_files) {

        // 读取预测图片
        byte[] bytes = Files.readAllBytes(Paths.get(f.getPath()));

        try (Tensor<String> input = Tensors.create(bytes);
             Tensor<Float> output =
                     session
                             .runner()
                             .feed("DecodeJpeg/contents:0", input)
                             .fetch("final_result:0")
                             .run()
                             .get(0)
                             .expect(Float.class)) {
          predictionData = new float[(int) output.shape()[0]][(int) output.shape()[1]];
          output.copyTo(predictionData);
          printPredictionData(f.getName(), predictionData);
          predictionData = null;
        }
      }
    }
  }

  private static byte[] loadTensorGraph(String graphPath) throws IOException {
    return Files.readAllBytes(Paths.get(graphPath));
  }

  private static List<String> loadLabels(String labelPath) throws IOException {
    return Files.readAllLines(Paths.get(labelPath), Charset.forName("utf-8"));
  }

  private static void printPredictionData(String filename, float[][] predictionData) {
    System.out.printf("\nfile : %s\n", filename);
    for (int i = 0; i < predictionData.length; i++) {
      for (int j = 0; j < predictionData[i].length; j++) {
        System.out.printf("label %s : %.2f%% likely\n", (j + 1), predictionData[i][j] * 100);
      }
    }
  }

  private static int argmax(float[] predictionData) {
    int best = 0;
    for (int i = 1; i < predictionData.length; ++i) {
      if (predictionData[i] > predictionData[best]) {
        best = i;
      }
    }
    return best;
  }

  private static void printAllOperation(Graph graph) {
    Iterator<Operation> it = graph.operations();
    Operation op = null;
    while (it.hasNext()) {
      op = it.next();
      System.out.printf("option name: %s\n", op.name());
    }
  }
}
